#include "tsb.h"
#include <iostream>
#include <stdexcept>
#include <cstring>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <mutex>
#include "selutility.h"

namespace TSBLib {

    // 构造函数：初始化连接
    TSB::TSB() : sock(-1) {
        Connect();
    }

    // 析构函数：关闭连接
    TSB::~TSB() {
        Disconnect();
    }

    // 建立与swtpm的连接
    void TSB::Connect() {
        const std::string swtpm_ip = "127.0.0.1";
        const int swtpm_port = 2321;

        // 创建套接字
        if ((sock = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
            throw std::runtime_error("Socket creation error");
        }

        struct sockaddr_in serv_addr;
        serv_addr.sin_family = AF_INET;
        serv_addr.sin_port = htons(swtpm_port);

        // 将IP地址从文本转换为二进制形式
        if(inet_pton(AF_INET, swtpm_ip.c_str(), &serv_addr.sin_addr) <= 0) {
            close(sock);
            throw std::runtime_error("Invalid address/ Address not supported");
        }

        // 连接到swtpm服务器
        if (connect(sock, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) {
            close(sock);
            throw std::runtime_error("Connection Failed");
        }

        std::cout << "Connected to swtpm at " << swtpm_ip << ":" << swtpm_port << std::endl;
    }

    // 关闭连接
    void TSB::Disconnect() {
        if (sock != -1) {
            close(sock);
            sock = -1;
            std::cout << "Disconnected from swtpm." << std::endl;
        }
    }

    // 计算SHA256哈希值
    std::string TSB::ComputeSHA256(const void* data, size_t size) {
        SHA256 sha;
        sha.update(data, size);
        return sha.final();
    }

    // 将内存数据发送到swtpm并接收响应
    int TSB::ForwardMemory(const void* input_data, size_t input_size, void* output_data, size_t& output_size) {
        if (sock == -1) {
            throw std::runtime_error("Not connected to swtpm");
        }

        std::vector<uint8_t> _input_data((uint8_t *)input_data, (uint8_t *)(input_data) + input_size);
        selutility::addIpmiSelLog("/xyz/openbmc_project/sensors/tpm/tpmcmd", _input_data);

        // 发送数据
        size_t total_sent = 0;
        const uint8_t* buffer = static_cast<const uint8_t*>(input_data);
        while (total_sent < input_size) {
            ssize_t sent = send(sock, buffer + total_sent, input_size - total_sent, 0);
            if (sent == -1) {
                throw std::runtime_error("Failed to send data");
            }
            total_sent += sent;
        }

        // 接收响应数据
        size_t total_received = 0;
        uint8_t* out_buffer = static_cast<uint8_t*>(output_data);
        while (total_received < output_size) {
            ssize_t received = recv(sock, out_buffer + total_received, output_size - total_received, 0);
            if (received == -1) {
                throw std::runtime_error("Failed to receive data");
            }
            if (received == 0) { // 连接关闭
                break;
            }
            total_received += received;
        }

        std::vector<uint8_t> _output_data((uint8_t *)output_data, (uint8_t *)(output_data) + output_size);
        selutility::addIpmiSelLog("/xyz/openbmc_project/sensors/tpm/tpmres", _output_data);

        output_size = total_received;

        return 1;
    }

}
